import os
from collections import defaultdict
import pandas as pd
import pymysql
from export.db_config import db_config


class DataBaseQuery:
    # 连接数据库
    def execute_query(self, query, params=None):
        connection = pymysql.connect(**db_config)
        cursor = connection.cursor()
        cursor.execute(query, params)
        results = cursor.fetchall()
        columns = [desc[0] for desc in cursor.description]
        cursor.close()
        connection.close()
        return pd.DataFrame(results, columns=columns)

    # 获取所有区域站点配置
    def query_area_list(self):
        return self.execute_query("SELECT * FROM area_recruit_config")

    # 获取被试的生日
    def query_subject_birthday_list(self):
        return self.execute_query('''
           select subject_id, birthday
           from pers_subject
           where project_id='hor'
           ''')

    # 查询被试的模态采集数据
    def query_stat_clct_subject_list(self):
        return self.execute_query("SELECT * FROM statistics_clct_subject")

    # 获取被试模态第一次采集时间
    def query_modality_first_clct_list(self):
        return self.execute_query('''
        select site_id, subject_id, min(start_clct_date) minClctDate
        from statistics_clct_subject
        where project_id = 'hor'
          and start_clct_date is not null
          and modality_id not in ('env')
        group by site_id, subject_id
        ''')

    #获取被试的家长数据
    def query_stat_env_clct_subject_list(self):
        return self.execute_query("SELECT * FROM statistics_clct_nsubject where survey_type = 1")


class SubjectStatExport:
    def __init__(self):
        self.db_query = DataBaseQuery()

    #导出
    def subject_group_by_age_stat_export(self):
        project_id = "HOR"

        head_list = self.build_head_list()
        data_list = self.build_data_list(project_id)

        self.export_excel_two_month_site(head_list, data_list, "subjectGroupByAgeStat")

    #构建data数据
    def build_data_list(self, project_id):

        # 获取站点区域数据
        area_list = self.db_query.query_area_list()
        area_list = area_list[area_list['project_id'].str.upper() == project_id]

        # 根据站点分组
        site_id_to_stat_map = self.db_query.query_stat_clct_subject_list().groupby('site_id')
        site_id_to_stat_env_map = self.db_query.query_stat_env_clct_subject_list().groupby('site_id')
        # map，根据siteId分组，在根据被试的采集年龄分组
        site_id_age_to_subject_id_list_map = self.get_site_id_age_to_subject_id_list_map()

        data_list = []
        for row in area_list.itertuples():
            site_id = row.site_id
            site_data_list = [row.site_name]

            # map, 获取站点对应的age对应的被试
            age_to_subject_id_list_map = site_id_age_to_subject_id_list_map.get(site_id, defaultdict(list))

            #获取站点下对应采集数据
            stat_list = site_id_to_stat_map.get_group(site_id)
            stat_env_list = site_id_to_stat_env_map.get_group(site_id)

            # 三三组和，4种，算上采集完成，8种
            three_list = [["mri", "eeg", "bhv"], ["mri", "eeg", "bio"], ["mri", "bhv", "bio"], ["eeg", "bhv", "bio"]]
            for age in range(6, 19):

                # 被试id list
                subject_id_list = age_to_subject_id_list_map.get(age, [])

                # 过滤出指定年龄的采集数据，并根据subject_id分组
                # 列表推导式，返回的是一个普通list
                # age_stat_list = [row for index, row in stat_list.iterrows() if row['subject_id'] in subject_id_list]
                # print(type(age_stat_list))
                # #query 方法允许你使用类似于 SQL 的查询语句来筛选 DataFrame。@ 符号用于引用外部变量 subject_id_list。
                # age_stat_list = stat_list.query('subject_id in @subject_id_list')
                # print(age_stat_list['subject_id'])
                # 使用isin判断是否包含
                age_stat_list = stat_list[stat_list['subject_id'].isin(subject_id_list)]
                # 进行分组，根据subjectId
                subjectId_statList_map = age_stat_list.groupby('subject_id')

                age_stat_env_List = stat_env_list.query('subject_id in @subject_id_list')
                subjectId_statEnvList_map = age_stat_env_List.groupby('subject_id')

                for modality_list in three_list:
                    # 获取采集的被试数
                    clct_num_subject_id_set = self.get_clct_num(subjectId_statList_map, modality_list)
                    com_num = self.get_com_num(subjectId_statList_map, subjectId_statEnvList_map, modality_list,
                                               clct_num_subject_id_set)
                    site_data_list.append(len(clct_num_subject_id_set))
                    site_data_list.append(com_num)

            # 四组和，1种，算上采集完成，2种
            four_list = ["mri", "eeg", "bhv", "bio"]
            for age in range(6, 19):
                # 被试id list
                subject_id_list = age_to_subject_id_list_map.get(age, [])

                # 过滤出指定年龄的采集数据，并根据subject_id分组
                # 使用isin判断是否包含
                age_stat_list = stat_list[stat_list['subject_id'].isin(subject_id_list)]
                # 进行分组，根据subjectId
                subjectId_statList_map = age_stat_list.groupby('subject_id')

                age_stat_env_List = stat_env_list.query('subject_id in @subject_id_list')
                subjectId_statEnvList_map = age_stat_env_List.groupby('subject_id')

                clct_num_subject_id_set = self.get_clct_num(subjectId_statList_map, four_list)
                com_num = self.get_com_num(subjectId_statList_map, subjectId_statEnvList_map, four_list,
                                           clct_num_subject_id_set)
                site_data_list.append(len(clct_num_subject_id_set))
                site_data_list.append(com_num)

            data_list.append(site_data_list)
        return data_list

    # 根据siteId分组，在根据被试的采集年龄分组
    def get_site_id_age_to_subject_id_list_map(self):
        site_id_subject_id_to_first_clct_map = self.db_query.query_modality_first_clct_list().groupby(
            ['site_id', 'subject_id'])
        subject_id_to_birthday_map = self.db_query.query_subject_birthday_list().groupby('subject_id')

        site_id_age_to_subject_id_list_map = {}
        for site_id_subject_id, subject_stat in site_id_subject_id_to_first_clct_map:
            subjectId = site_id_subject_id[1]
            siteId = site_id_subject_id[0]

            first_clct = subject_stat['minClctDate'].iloc[0]
            birth_date = subject_id_to_birthday_map.get_group(subjectId)['birthday'].iloc[0]
            age = (first_clct - birth_date).days // 365

            # 确保 siteId 对应的内层字典存在
            if siteId not in site_id_age_to_subject_id_list_map:
                site_id_age_to_subject_id_list_map[siteId] = {}

            # 确保 age 对应的列表存在
            if age not in site_id_age_to_subject_id_list_map[siteId]:
                site_id_age_to_subject_id_list_map[siteId][age] = []

            site_id_age_to_subject_id_list_map[siteId][age].append(subjectId)
        return site_id_age_to_subject_id_list_map

    # 获取完成数，且需要在已采集的被试基础上
    def get_com_num(self, subjectId_statList_map, subjectId_statEnvList_map, modality_list, clct_num_subject_id_set):
        com_num = 0
        for subject_id, stats in subjectId_statList_map:
            if subject_id not in clct_num_subject_id_set:
                continue

            # 过滤出完成时间不为null的
            com_list = stats[stats['complete_clct_date'].notnull()]
            if len(com_list) == len(modality_list):
                if all(row.modality_id in modality_list for index, row in com_list.iterrows()):
                    if 'bhv' in modality_list:
                        if subject_id in subjectId_statEnvList_map.groups:
                            com_date = subjectId_statEnvList_map.get_group(subject_id)['complete_clct_date'].iloc[0]
                            if com_date is not None:
                                com_num += 1
        return com_num

    # 获取采集的被试数，被试采集的模态数符合，且模态一一对应
    def get_clct_num(self, subjectId_statList_map, modality_list):
        result_set = set()
        for subject_id, stats in subjectId_statList_map:
            if len(stats) == len(modality_list):
                if all(row.modality_id in modality_list for index, row in stats.iterrows()):
                    result_set.add(subject_id)
        return result_set

    # 构建head数据
    def build_head_list(self):
        head_list = [["站点名称"]]

        # 三三组和，4种，算上采集完成，8种
        three_list = ["磁共振&脑电&认知行为与环境", "磁共振&脑电&基因", "磁共振&认知行为与环境&基因",
                      "脑电&认知行为与环境&基因"]
        for age in range(6, 19):
            for three in three_list:
                head_list.append([f"{age}岁仅采集{three}"])
                head_list.append([f"{age}岁仅完成{three}"])

        # 四组和，1种，算上采集完成，2种
        for age in range(6, 19):
            head_list.append([f"{age}岁仅采集磁共振&脑电&认知行为与环境&基因"])
            head_list.append([f"{age}岁仅完成磁共振&脑电&认知行为与环境&基因"])

        return head_list

    def export_excel_two_month_site(self, head_list, data_list, file_name):
        df = pd.DataFrame(data_list, columns=[item[0] for item in head_list])
        file_path = os.path.join("D:\\", f"{file_name}.xlsx")
        os.makedirs(os.path.dirname(file_path), exist_ok=True)
        df.to_excel(file_path, index=False, sheet_name=file_name)

        print(f"文件导出成功: {file_path}")


if __name__ == "__main__":
    exporter = SubjectStatExport()
    exporter.subject_group_by_age_stat_export()
